
    #include <mpi.h>
    #include <stdio.h>
    #include <stdlib.h>
    
    #ifndef _EXTERN_C_
    #ifdef __cplusplus
    #define _EXTERN_C_ extern "C"
    #else /* __cplusplus */
    #define _EXTERN_C_
    #endif /* __cplusplus */
    #endif /* _EXTERN_C_ */
    
    #ifdef MPICH_HAS_C2F
    _EXTERN_C_ void *MPIR_ToPointer(int);
    #endif // MPICH_HAS_C2F
    
    #ifdef PIC
    /* For shared libraries, declare these weak and figure out which one was linked
    based on which init wrapper was called.  See mpi_init wrappers.  */
    #pragma weak pmpi_init
    #pragma weak PMPI_INIT
    #pragma weak pmpi_init_
    #pragma weak pmpi_init__
    #endif /* PIC */
    
    _EXTERN_C_ void pmpi_init(MPI_Fint *ierr);
    _EXTERN_C_ void PMPI_INIT(MPI_Fint *ierr);
    _EXTERN_C_ void pmpi_init_(MPI_Fint *ierr);
    _EXTERN_C_ void pmpi_init__(MPI_Fint *ierr);
    
    /* Customized MPI Wrapper */ 

#include "array.h"
#define VERTICES_TAG 44
#define GRAPH_TAG 123
#define ROOT 0
#define SUCCESS 0
#define ERROR -1
#define OPER_COUNTER(op, name) {name, 0}


static int initialized = 0;
static double start_time = 0.0;
static double end_time = 0.0;
static double in_time = 0.0;
static double out_time = 0.0;
static const char* dot_fname = "MPI_task_graph.dot";
static const char* pdf_fname = "MPI_task_graph.pdf";
static const char* critical_fname = "critPath.out";
static const char* stat_fname = "stats.dat";
static int vid = 0;

typedef enum Mpi_type_t{
    MPI_Barrier_counter, 
    MPI_Alltoall_counter, 
    MPI_Scatter_counter, 
    MPI_Gather_counter, 
    MPI_Reduce_counter, 
    MPI_Allreduce_counter, 
    MPI_Send_counter, 
    MPI_Recv_counter, 
    MPI_Isend_counter, 
    MPI_Irecv_counter, 
    MPI_Wait_counter, 
    MPI_Waitall_counter, 
    MPI_Init_counter, 
    MPI_Finalize_counter, 
    MPI_Total_counter

} MPI_ID;

typedef enum Edge_type_t{
    COMPUTATION_EDGE,
    LATENCY_EDGE,
} EDGE_TYPE;

typedef struct Operation_counter {
    const char *name;
    unsigned long long value;
} Operation_counter_t;

Operation_counter_t counters[] = {
    OPER_COUNTER(MPI_Barrier_counter, "MPI_Barrier"),
    OPER_COUNTER(MPI_Alltoall_counter, "MPI_Alltoall"),
    OPER_COUNTER(MPI_Scatter_counter, "MPI_Scatter"),
    OPER_COUNTER(MPI_Gather_counter, "MPI_Gather"),
    OPER_COUNTER(MPI_Reduce_counter, "MPI_Reduce"),
    OPER_COUNTER(MPI_Allreduce_counter, "MPI_Allreduce"),
    OPER_COUNTER(MPI_Send_counter, "MPI_Send"),
    OPER_COUNTER(MPI_Recv_counter, "MPI_Recv"),
    OPER_COUNTER(MPI_Isend_counter, "MPI_Isend"),
    OPER_COUNTER(MPI_Irecv_counter, "MPI_Irecv"),
    OPER_COUNTER(MPI_Wait_counter, "MPI_Wait"),
    OPER_COUNTER(MPI_Waitall_counter, "MPI_Waitall"),
    OPER_COUNTER(MPI_Init_counter, "MPI_Init"),
    OPER_COUNTER(MPI_Finalize_counter, "MPI_Finalize"),

};

typedef struct MPI_Request_node_t MPI_Request_node;
void Pre_MPI(void);
int Post_MPI(MPI_ID id);

// Edge list of a graph
typedef struct Edge_t Edge;
typedef struct Vertex_t Vertex; 
struct Edge_t
{
    double weight;
    Vertex* dest_vertex;
    Edge *edge_ptr;
    int on_crit_path;
    EDGE_TYPE type;
};

// Vertex list of a graph
struct Vertex_t
{
    int out_degree;
    int in_degree;
    long long int vertex_id;
    int visited;
    Edge *edge_ptr;
    Vertex *next;
    Vertex *prev_crit_vertex;
    double crit_duration;
    int on_crit_path;
    // MPI information 
    int rank;
    int src_rank;
    int dest_rank;
    double weight;
    MPI_ID mpi_id;
    long long unsigned msg_size;
    //added by huangfq, for graph merging
    long long int wait_vertex_id;
    MPI_Datatype msg_type;
    int is_matched; //for matching purpose only
    int is_collective; //determine if the operation is collective opertion
    double in_time; //the start time of the operation
    double out_time; //the end time of the operation
    long long int invoke_id;//for each operation
};

// Graph is represented by a list of vertex
typedef struct Graph_t
{   
    int num_vertices;
    Vertex *vertex_head;
} Graph;

// Linked path represents a path of a graph
typedef struct Node_t Node;

// A node of a path (represented by a vertex)
struct Node_t
{
    Vertex *vertex;
    Node *next;
};

typedef struct Path_t
{
    Node *head;
} Path;


// A linked list handle MPI_Request
struct MPI_Request_node_t
{
    Vertex *vertex;
    MPI_Request* request;
    MPI_Request_node* next;
    int is_matched;
};

MPI_Request_node *request_head = NULL;


////////////////////////////////////////////////////////////
////                Graph Operations                    ////
////////////////////////////////////////////////////////////

Graph* create_graph();

// Initialize a vertex to default values with id
void init_vertex
(
    Vertex *v, 
    long long unsigned vid, 
    int rank,
    double in_time,
    double out_time,
    double exec_time,
    MPI_ID id,
    long long unsigned invoke_id,
    int is_collective
);

// Insert a vertex to graph
int insert_vertex(Graph *graph, Vertex* vertex);

// Get the last vertex of graph
Vertex* get_last_vertex(const Graph* graph);

// Insert an edge (vertex1, vertex2) to graph
int insert_edge
(
    Vertex *vertex1, 
    Vertex *vertex2, 
    double weight,
    EDGE_TYPE type 
);

int insert_computation_edge(Vertex *v1, Vertex *V2);
int insert_latency_edge(Vertex *v1, Vertex *v2);
double get_collective_weight(Vertex *v, int num_nodes);
Graph* merge_graph(Graph *subgraphs[], int num_nodes);
void merge_collective_vertices(Graph *subgraphs[], int num_nodes);
void add_latency_edges(Graph *subgraphs[], int num_nodes);
void add_computation_edges(Graph *subgraphs[], int num_nodes);
long long int get_actual_msg_size(MPI_Datatype type, long long int count);

// Find the longest path (Critical path) of a graph
int compute_critical_path
(
    const Path* topo_path, 
    Path* crit_path,
    double *length
);

// Sort the vertices in topological order to path
void topological_sort
(
    const Graph *graph, 
    Path *path
);

// Helper function for topological sort
void visit(Vertex *vetex, Path *path);

// Print the graph
void print_graph(Graph *graph);

// Write graph to dot format and creates a pdf file
int generate_graph
(
    const Graph *graph, 
    const char* dotName, 
    const char* pdfName
);

// Output the critical path into a file
int output_crit_path
(
    const Path* crit_path,
    const Graph* graph,
    const char* fname
);

void free_graph(Graph *graph);

/* \/ \/ \/ \/ \/ \/ \/ \/ \/ \/ \/ \/ \/ \/ \/ \/ \/ \/ \/ \/ \/ \/ \/ \/ \/ */

// Insert a vertex to the beginning of a path
void insert(Path *path, Vertex* vertex);

// Append a vertex to the end of a path
void append(Path *path, Vertex* vertex);

// Print path
void print_path(const Path *path);

void Pre_MPI(void);
int Post_MPI(MPI_ID id);

// Display the median, average, min and max of MPI operation execution time 
void output_mpi_statistics
(
     const char* fname, 
     Graph* graph,
     int num_graphs
 );

Graph * per_rank_graph = NULL;
Graph ** root_graph = NULL;
int *num_vertices = NULL;

///////////////////////////////////////////////////////////////////////////////
//                               Add generic wrapper                         //
///////////////////////////////////////////////////////////////////////////////

/* ================== C Wrappers for MPI_Barrier ================== */
_EXTERN_C_ int PMPI_Barrier(MPI_Comm comm);
_EXTERN_C_ int MPI_Barrier(MPI_Comm comm);

/* ================== C Wrappers for MPI_Alltoall ================== */
_EXTERN_C_ int PMPI_Alltoall(void *sendbuf, int sendcount, MPI_Datatype sendtype, void *recvbuf, int recvcount, MPI_Datatype recvtype, MPI_Comm comm);
_EXTERN_C_ int MPI_Alltoall(void *sendbuf, int sendcount, MPI_Datatype sendtype, void *recvbuf, int recvcount, MPI_Datatype recvtype, MPI_Comm comm);

/* ================== C Wrappers for MPI_Scatter ================== */
_EXTERN_C_ int PMPI_Scatter(void *sendbuf, int sendcount, MPI_Datatype sendtype, void *recvbuf, int recvcount, MPI_Datatype recvtype, int root, MPI_Comm comm);
_EXTERN_C_ int MPI_Scatter(void *sendbuf, int sendcount, MPI_Datatype sendtype, void *recvbuf, int recvcount, MPI_Datatype recvtype, int root, MPI_Comm comm);

/* ================== C Wrappers for MPI_Gather ================== */
_EXTERN_C_ int PMPI_Gather(void *sendbuf, int sendcount, MPI_Datatype sendtype, void *recvbuf, int recvcount, MPI_Datatype recvtype, int root, MPI_Comm comm);
_EXTERN_C_ int MPI_Gather(void *sendbuf, int sendcount, MPI_Datatype sendtype, void *recvbuf, int recvcount, MPI_Datatype recvtype, int root, MPI_Comm comm);

/* ================== C Wrappers for MPI_Reduce ================== */
_EXTERN_C_ int PMPI_Reduce(void *sendbuf, void *recvbuf, int count, MPI_Datatype datatype, MPI_Op op, int root, MPI_Comm comm);
_EXTERN_C_ int MPI_Reduce(void *sendbuf, void *recvbuf, int count, MPI_Datatype datatype, MPI_Op op, int root, MPI_Comm comm);

/* ================== C Wrappers for MPI_Allreduce ================== */
_EXTERN_C_ int PMPI_Allreduce(void *sendbuf, void *recvbuf, int count, MPI_Datatype datatype, MPI_Op op, MPI_Comm comm);
_EXTERN_C_ int MPI_Allreduce(void *sendbuf, void *recvbuf, int count, MPI_Datatype datatype, MPI_Op op, MPI_Comm comm);

/* ================== C Wrappers for MPI_Send ================== */
_EXTERN_C_ int PMPI_Send(void *buf, int count, MPI_Datatype datatype, int dest, int tag, MPI_Comm comm);
_EXTERN_C_ int MPI_Send(void *buf, int count, MPI_Datatype datatype, int dest, int tag, MPI_Comm comm);

/* ================== C Wrappers for MPI_Recv ================== */
_EXTERN_C_ int PMPI_Recv(void *buf, int count, MPI_Datatype datatype, int source, int tag, MPI_Comm comm, MPI_Status *status);
_EXTERN_C_ int MPI_Recv(void *buf, int count, MPI_Datatype datatype, int source, int tag, MPI_Comm comm, MPI_Status *status);

/* ================== C Wrappers for MPI_Isend ================== */
_EXTERN_C_ int PMPI_Isend(void *buf, int count, MPI_Datatype datatype, int dest, int tag, MPI_Comm comm, MPI_Request *request);
_EXTERN_C_ int MPI_Isend(void *buf, int count, MPI_Datatype datatype, int dest, int tag, MPI_Comm comm, MPI_Request *request);

/* ================== C Wrappers for MPI_Irecv ================== */
_EXTERN_C_ int PMPI_Irecv(void *buf, int count, MPI_Datatype datatype, int source, int tag, MPI_Comm comm, MPI_Request *request);
_EXTERN_C_ int MPI_Irecv(void *buf, int count, MPI_Datatype datatype, int source, int tag, MPI_Comm comm, MPI_Request *request);

/* ================== C Wrappers for MPI_Wait ================== */
_EXTERN_C_ int PMPI_Wait(MPI_Request *request, MPI_Status *status);
_EXTERN_C_ int MPI_Wait(MPI_Request *request, MPI_Status *status);

/* ================== C Wrappers for MPI_Waitall ================== */
_EXTERN_C_ int PMPI_Waitall(int count, MPI_Request *array_of_requests, MPI_Status *array_of_statuses);
_EXTERN_C_ int MPI_Waitall(int count, MPI_Request *array_of_requests, MPI_Status *array_of_statuses);

/* ================== C Wrappers for MPI_Init ================== */
_EXTERN_C_ int PMPI_Init(int *argc, char ***argv);
_EXTERN_C_ int MPI_Init(int *argc, char ***argv);

/* ================== C Wrappers for MPI_Finalize ================== */
_EXTERN_C_ int PMPI_Finalize();
_EXTERN_C_ int MPI_Finalize();


